# Copyright (c) 2025 Huawei Technologies Co., Ltd.
# openUBMC is licensed under Mulan PSL v2.
# You can use this software according to the terms and conditions of the Mulan PSL v2.
# You may obtain a copy of Mulan PSL v2 at:
#         http://license.coscl.org.cn/MulanPSL2
# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
# See the Mulan PSL v2 for more details.

# -*- coding: utf-8 -*-
"""
异步数据流分析测试

测试 dataflow_analysis.py 中的异步数据流分析功能：
- Segment 划分
- 最小代码单元划分
- 变量提升规则（规则1-3）
"""

from mcpy.compiler import MCPyCompiler

# 辅助函数定义（用于测试）
HELPER_FUNCTIONS = """
async def async_step1() -> int:
    return 1

async def async_step2() -> int:
    return 2

async def async_step3() -> int:
    return 3

async def async_step4() -> int:
    return 4

async def async_step(i: int) -> int:
    return i

async def get_condition() -> bool:
    return True

async def step1() -> int:
    return 1

async def step2() -> int:
    return 2

async def step3() -> int:
    return 3

async def step4() -> int:
    return 4
"""


def compile_and_get_async_dataflow(source_code: str, add_helpers: bool = True):
    """编译源代码并获取异步数据流分析结果"""
    # 自动添加辅助函数
    if add_helpers:
        source_code = HELPER_FUNCTIONS + "\n" + source_code

    compiler = MCPyCompiler()
    result = compiler.compile_source(source_code)
    ir_module = result.ir_module

    # 检查编译是否成功
    if ir_module is None:
        print(f"编译失败！")
        if hasattr(result, "errors") and result.errors:
            print(f"错误信息: {result.errors}")
        # 如果是测试用例，打印更详细的错误信息
        if hasattr(result, "error_messages") and result.error_messages:
            print(f"详细错误: {result.error_messages}")
        raise RuntimeError("编译失败")

    # 收集所有异步函数及其数据流分析结果
    async_functions = {}
    for node in ir_module.walk():
        from mcpy.ir.declarations.function import IRFunctionNode

        if isinstance(node, IRFunctionNode) and node.is_async:
            async_dataflow = node.get_extension_param("async_dataflow")
            if async_dataflow:
                async_functions[node.name] = (node, async_dataflow)

    return ir_module, async_functions


# 辅助函数：验证 Segment
def assert_segment(segment, segment_id, start, end, has_await):
    """验证 Segment 的属性"""
    assert segment.segment_id == segment_id
    assert segment.start_stmt_index == start
    assert segment.end_stmt_index == end
    assert segment.has_await == has_await


# 辅助函数：验证最小代码单元
def assert_unit(unit, unit_id, start, end, has_await, vars_defined):
    """验证最小代码单元的属性"""
    assert unit.unit_id == unit_id
    assert unit.start_stmt_index == start
    assert unit.end_stmt_index == end
    assert unit.has_await == has_await

    # 验证变量定义
    for var in vars_defined:
        assert var in unit.variables_defined

    # 验证没有多余的变量
    extra_vars = set(unit.variables_defined) - set(vars_defined)
    if extra_vars:
        # 允许外部函数名（如 async_step1）
        external_funcs = {
            v
            for v in extra_vars
            if v.startswith("async_") or v.startswith("get_") or v.startswith("step")
        }
        extra_vars -= external_funcs
        assert not extra_vars


# 辅助函数：验证变量提升
def assert_var_promotion(
    async_dataflow, var_name, promotion_type, defined_in_unit=None
):
    """验证变量提升类型"""
    assert var_name in async_dataflow.variable_promotions

    promotion = async_dataflow.variable_promotions[var_name]
    assert promotion.promotion_type == promotion_type

    if defined_in_unit is not None:
        assert promotion.defined_in_unit_id == defined_in_unit

    # 验证变量在相应的列表中
    if promotion_type == "member":
        assert var_name in async_dataflow.member_variables
        assert var_name not in async_dataflow.header_variables
        # local_variables 是字典，键是 Unit ID，值是变量列表
        assert not any(
            var_name in vars for vars in async_dataflow.local_variables.values()
        )
    elif promotion_type == "header":
        assert var_name in async_dataflow.header_variables
        assert var_name not in async_dataflow.member_variables
        # local_variables 是字典，键是 Unit ID，值是变量列表
        assert not any(
            var_name in vars for vars in async_dataflow.local_variables.values()
        )
    elif promotion_type == "local":
        # local_variables 是字典，键是 Unit ID，值是变量列表
        assert any(var_name in vars for vars in async_dataflow.local_variables.values())
        assert var_name not in async_dataflow.member_variables
        assert var_name not in async_dataflow.header_variables


def test_simple_await_segment():
    """测试简单的 await 语句 Segment 划分"""

    code = """
async def simple_await() -> int:
    x = await async_step1()
    y = await async_step2()
    return x + y
"""

    ir_module, async_funcs = compile_and_get_async_dataflow(code)
    assert "simple_await" in async_funcs

    func_node, async_dataflow = async_funcs["simple_await"]

    # 验证 Segment 划分（3 个 Segment）
    assert len(async_dataflow.segments) == 3
    assert_segment(async_dataflow.segments[0], 0, 0, 1, has_await=True)
    assert_segment(async_dataflow.segments[1], 1, 1, 2, has_await=True)
    assert_segment(async_dataflow.segments[2], 2, 2, 3, has_await=False)

    # 验证最小代码单元划分（3 个单元，与 Segment 重叠）
    # await 结果变量的定义属于下一个单元
    assert len(async_dataflow.units) == 3
    assert_unit(async_dataflow.units[0], 0, 0, 1, has_await=True, vars_defined=[])
    assert_unit(async_dataflow.units[1], 1, 1, 2, has_await=True, vars_defined=["x"])
    assert_unit(async_dataflow.units[2], 2, 2, 3, has_await=False, vars_defined=["y"])

    # 验证变量提升
    assert_var_promotion(async_dataflow, "x", "member", defined_in_unit=1)
    assert_var_promotion(async_dataflow, "y", "local", defined_in_unit=2)


def test_await_in_if_branch():
    """测试 if 分支中的 await - 综合测试规则1和规则3
    Python 的 if/else 不创建作用域，两个分支的 result 是同一个变量。
    变量提升采用全局分析：收集所有使用位置，只要有任何使用跨 await 就提升为 member。
    """

    code = """
async def await_in_if(condition: bool) -> int:
    x = 0
    then_val = 10
    else_val = 20
    extra_val = 30
    mixed_val = 40  # 在 then 分支 await 之前使用，在 else 分支 await 之后使用
    if condition:
        then_val = then_val + mixed_val  # await 之前使用 mixed_val
        result = await async_step()
        x = result + extra_val
    else:
        result = await async_step(else_val)
        x = result + extra_val + mixed_val  # await 之后使用 mixed_val
    return x
"""

    ir_module, async_funcs = compile_and_get_async_dataflow(code)
    assert "await_in_if" in async_funcs

    func_node, async_dataflow = async_funcs["await_in_if"]

    # 验证 Segment 划分（6 个 Segment）
    # 外层：
    #   Segment 0: [0, 6) - x=0, then_val=10, else_val=20, extra_val=30, mixed_val=40, if 语句（包含子语句块，有 await）
    #   Segment 5: [6, 7) - return x
    # then 分支：
    #   Segment 1: [0, 2) - then_val = then_val + mixed_val, await async_step()
    #   Segment 2: [2, 3) - x = result + extra_val
    # else 分支：
    #   Segment 3: [0, 1) - await async_step(else_val)
    #   Segment 4: [1, 2) - x = result + extra_val + mixed_val
    assert len(async_dataflow.segments) == 6

    assert_segment(async_dataflow.segments[0], 0, 0, 6, has_await=True)
    assert async_dataflow.segments[0].contains_sub_blocks is True
    assert_segment(async_dataflow.segments[1], 1, 0, 2, has_await=True)
    assert async_dataflow.segments[1].contains_sub_blocks is False
    assert_segment(async_dataflow.segments[2], 2, 2, 3, has_await=False)
    assert async_dataflow.segments[2].contains_sub_blocks is False
    assert_segment(async_dataflow.segments[3], 3, 0, 1, has_await=True)
    assert async_dataflow.segments[3].contains_sub_blocks is False
    assert_segment(async_dataflow.segments[4], 4, 1, 2, has_await=False)
    assert async_dataflow.segments[4].contains_sub_blocks is False
    assert_segment(async_dataflow.segments[5], 5, 6, 7, has_await=False)
    assert async_dataflow.segments[5].contains_sub_blocks is False

    # 验证最小代码单元划分（6 个单元）
    # 外层：
    #   Unit 0: [0, 5) - x=0, then_val=10, else_val=20, extra_val=30, mixed_val=40（if 之前）
    #   Unit 5: [6, 7) - return x（if 之后）
    # then 分支：
    #   Unit 1: [0, 2) - then_val = then_val + mixed_val, await async_step()，vars=[]，parent=0
    #   Unit 2: [2, 3) - x = result + extra_val，vars=['result']，parent=0
    # else 分支：
    #   Unit 3: [0, 1) - await async_step(else_val)，vars=[]，parent=0
    #   Unit 4: [1, 2) - x = result + extra_val + mixed_val，vars=[]，parent=0
    assert len(async_dataflow.units) == 6

    # 验证外层主语句块的单元（if 之前）
    assert_unit(
        async_dataflow.units[0],
        0,
        0,
        5,
        has_await=False,
        vars_defined=["x", "then_val", "else_val", "extra_val", "mixed_val"],
    )

    # 验证 then 分支的单元
    # Unit 1: [0, 2) 包含 `then_val = then_val + 1` 和 `result = await async_step()`
    # `then_val = ...` 是赋值（is_declaration=False），不在 vars_defined 中
    assert_unit(
        async_dataflow.units[1],
        1,
        0,
        2,
        has_await=True,
        vars_defined=[],  # 只包含声明，`then_val` 是赋值
    )
    # Unit 2: [2, 3) 包含 `x = result + extra_val`
    # `result` 是 await 结果变量（在 Unit 1 的 await 之后定义）
    # `x = ...` 是赋值（is_declaration=False），不在 vars_defined 中
    assert_unit(
        async_dataflow.units[2],
        2,
        2,
        3,
        has_await=False,
        vars_defined=["result"],  # 只包含 await 结果变量
    )

    # 验证 else 分支的单元
    # Unit 3: [0, 1) 包含 `result = await async_step(else_val)`
    # `result = ...` 是赋值（is_declaration=False），不在 vars_defined 中
    assert_unit(async_dataflow.units[3], 3, 0, 1, has_await=True, vars_defined=[])
    # Unit 4: [1, 2) 包含 `x = result + extra_val`
    # `result` 和 `x` 都是赋值（is_declaration=False），不在 vars_defined 中
    assert_unit(
        async_dataflow.units[4],
        4,
        1,
        2,
        has_await=False,
        vars_defined=[],  # 都是赋值，不是声明
    )

    # 验证外层主语句块的单元（if 之后）
    assert_unit(async_dataflow.units[5], 5, 6, 7, has_await=False, vars_defined=[])

    # 验证变量提升
    # x: 在单元 0 定义，在单元 5 使用（return x）
    #    注意：Unit 2 和 Unit 4 中的 `x = ...` 是赋值（写入），不是使用（读取）
    #    虽然 Unit 0 和 Unit 5 在同一语句块，但中间有 Unit 1, 3 包含 await
    #    规则1：跨最小代码单元使用（跨 await）→ 成员变量
    assert_var_promotion(async_dataflow, "x", "member", defined_in_unit=0)

    # then_val: 在单元 0（外层，无 parent）定义，在单元 1（then 分支第一个单元，parent=0）使用
    #           规则3：在子语句块的第一个单元使用但不在该单元定义 → 函数头变量
    assert_var_promotion(async_dataflow, "then_val", "header", defined_in_unit=0)

    # else_val: 在单元 0（外层，无 parent）定义，在单元 3（else 分支第一个单元，parent=0）使用
    #           规则3：在子语句块的第一个单元使用但不在该单元定义 → 函数头变量
    assert_var_promotion(async_dataflow, "else_val", "header", defined_in_unit=0)

    # extra_val: 在单元 0（外层，无 parent）定义，在单元 2, 4（两个分支的第二个单元，parent=0）使用
    #            规则1：跨最小代码单元使用（跨 await），且不是第一个单元 → 成员变量
    assert_var_promotion(async_dataflow, "extra_val", "member", defined_in_unit=0)

    # mixed_val: 在单元 0 定义，在单元 1, 4 使用
    #            Unit 1: then 分支第一个单元，await 之前使用（满足规则3条件）
    #            Unit 4: else 分支第二个单元，await 之后使用（满足规则1条件）
    #            规则1 优先于规则3 → 成员变量
    assert_var_promotion(async_dataflow, "mixed_val", "member", defined_in_unit=0)

    # result: 在 Unit 2 定义，在 Unit 2, 4 使用
    #         根据 python 的规则，Unit 4 的 `result = ...` 是赋值，不是定义
    #         规则1：跨最小代码单元使用（跨 await）→ 成员变量
    #
    #         说明：虽然 C++ 的 if/else 会创建作用域，但我们遵循 Python 语义，
    #         将 result 视为函数级别的变量。代码生成时会提升为成员变量。
    assert_var_promotion(async_dataflow, "result", "member", defined_in_unit=2)


def test_local_variable_in_same_unit():
    """测试同一最小代码单元内的局部变量 - 规则2（不提升）"""

    code = """
async def local_in_unit() -> int:
    temp = 10
    x = await async_step1()
    result = x + temp
    return result
"""

    ir_module, async_funcs = compile_and_get_async_dataflow(code)
    assert "local_in_unit" in async_funcs

    func_node, async_dataflow = async_funcs["local_in_unit"]

    # 验证 Segment 划分（2 个 Segment）
    assert len(async_dataflow.segments) == 2
    assert_segment(async_dataflow.segments[0], 0, 0, 2, has_await=True)
    assert_segment(async_dataflow.segments[1], 1, 2, 4, has_await=False)

    # 验证最小代码单元划分（2 个单元，与 Segment 重叠）
    # await 结果变量 x 的定义属于下一个单元
    assert len(async_dataflow.units) == 2
    assert_unit(async_dataflow.units[0], 0, 0, 2, has_await=True, vars_defined=["temp"])
    assert_unit(
        async_dataflow.units[1],
        1,
        2,
        4,
        has_await=False,
        vars_defined=["x", "result"],
    )

    # 验证变量提升
    assert_var_promotion(async_dataflow, "temp", "member", defined_in_unit=0)
    assert_var_promotion(async_dataflow, "x", "local", defined_in_unit=1)
    assert_var_promotion(async_dataflow, "result", "local", defined_in_unit=1)


def test_cross_await_variable():
    """测试跨 await 使用的变量 - 规则1（成员变量）"""

    code = """
async def cross_await() -> int:
    x = 10
    y = await async_step1()
    z = x + y
    result = await async_step2(z)
    return result
"""

    ir_module, async_funcs = compile_and_get_async_dataflow(code)
    assert "cross_await" in async_funcs

    func_node, async_dataflow = async_funcs["cross_await"]

    # 验证 Segment 划分（3 个 Segment）
    # Segment 0: [0, 2) - x = 10 和 y = await async_step1()
    # Segment 1: [2, 4) - z = x + y 和 result = await async_step2(z)
    # Segment 2: [4, 5) - return result
    assert len(async_dataflow.segments) == 3
    assert_segment(async_dataflow.segments[0], 0, 0, 2, has_await=True)
    assert_segment(async_dataflow.segments[1], 1, 2, 4, has_await=True)
    assert_segment(async_dataflow.segments[2], 2, 4, 5, has_await=False)

    # 验证最小代码单元划分（3 个单元，与 Segment 重叠）
    # 单元 0: x = 10, await async_step1() 表达式，vars=['x']（y 属于下一个单元）
    # 单元 1: y = ..., z = x + y, await async_step2(z) 表达式，vars=['y', 'z']（result 属于下一个单元）
    # 单元 2: result = ..., return result，vars=['result']
    assert len(async_dataflow.units) == 3
    assert_unit(async_dataflow.units[0], 0, 0, 2, has_await=True, vars_defined=["x"])
    assert_unit(
        async_dataflow.units[1],
        1,
        2,
        4,
        has_await=True,
        vars_defined=["y", "z"],
    )
    assert_unit(
        async_dataflow.units[2],
        2,
        4,
        5,
        has_await=False,
        vars_defined=["result"],
    )

    # 验证变量提升
    # x: 在单元 0 定义，在单元 1 使用（z = x + y），跨越单元，应该是 member
    assert_var_promotion(async_dataflow, "x", "member", defined_in_unit=0)

    # y: 在单元 1 定义和使用（z = x + y），应该是 local
    assert_var_promotion(async_dataflow, "y", "local", defined_in_unit=1)

    # z: 在单元 1 定义和使用（await async_step2(z)），应该是 local
    assert_var_promotion(async_dataflow, "z", "local", defined_in_unit=1)

    # result: 在单元 2 定义和使用（return result），应该是 local
    assert_var_promotion(async_dataflow, "result", "local", defined_in_unit=2)


def test_await_in_loop():
    """测试循环中的 await"""

    code = """
async def await_in_loop() -> int:
    total = 0
    for i in range(3):
        value = await async_step(i)
        total = total + value
    return total
"""

    ir_module, async_funcs = compile_and_get_async_dataflow(code)
    assert "await_in_loop" in async_funcs

    func_node, async_dataflow = async_funcs["await_in_loop"]

    # 验证 Segment 划分（4 个 Segment）
    # 外层：
    #   Segment 0: [0, 2) - total = 0, for 循环（包含子语句块，有 await）
    #   Segment 3: [2, 3) - return total
    # 循环体：
    #   Segment 1: [0, 1) - await async_step(i)
    #   Segment 2: [1, 2) - value = ..., total = total + value
    assert len(async_dataflow.segments) == 4

    assert_segment(async_dataflow.segments[0], 0, 0, 2, has_await=True)
    assert async_dataflow.segments[0].contains_sub_blocks is True

    assert_segment(async_dataflow.segments[1], 1, 0, 1, has_await=True)
    assert async_dataflow.segments[1].contains_sub_blocks is False

    assert_segment(async_dataflow.segments[2], 2, 1, 2, has_await=False)
    assert async_dataflow.segments[2].contains_sub_blocks is False

    assert_segment(async_dataflow.segments[3], 3, 2, 3, has_await=False)
    assert async_dataflow.segments[3].contains_sub_blocks is False

    # 验证最小代码单元划分（4 个单元）
    # 外层：
    #   Unit 0: [0, 1) - total = 0（for 之前）
    #   Unit 3: [2, 3) - return total
    # 循环体：
    #   Unit 1: [0, 1) - await async_step(i)，parent=0
    #   Unit 2: [1, 2) - value = ..., total = total + value，parent=0
    assert len(async_dataflow.units) == 4

    assert_unit(
        async_dataflow.units[0],
        0,
        0,
        1,
        has_await=False,
        vars_defined=["total"],
    )
    assert_unit(async_dataflow.units[1], 1, 0, 1, has_await=True, vars_defined=[])
    assert_unit(
        async_dataflow.units[2],
        2,
        1,
        2,
        has_await=False,
        vars_defined=["value"],
    )
    assert_unit(async_dataflow.units[3], 3, 2, 3, has_await=False, vars_defined=[])

    # 验证变量提升
    # total: 在 Unit 0 定义，在 Unit 2, 3 使用，跨越 Unit 1 的 await → member
    assert_var_promotion(async_dataflow, "total", "member", defined_in_unit=0)

    # value: 在 Unit 2 定义和使用 → local
    assert_var_promotion(async_dataflow, "value", "local", defined_in_unit=2)

    # i: 循环变量，由 for 语句隐式定义
    #    在循环体中使用（Unit 1: async_step(i)）
    #    循环执行流程：初始化 i → 循环体（await） → co_yield → 继续 → 自增 i → 下一次迭代
    #    自增操作跨越了 await，所以 i 必须是 member
    assert_var_promotion(async_dataflow, "i", "member")


def test_await_before_if_with_await_in_branches():
    """测试 await 在 if 之前，if 分支中也有 await 的情况"""

    code = """
async def await_before_if() -> int:
    cond = await get_condition()
    x = 0
    if cond:
        x = await async_step1()
    else:
        x = await async_step2()
    return x
"""

    ir_module, async_funcs = compile_and_get_async_dataflow(code)
    assert "await_before_if" in async_funcs

    func_node, async_dataflow = async_funcs["await_before_if"]

    # 验证 Segment 划分（5 个 Segment）
    # 外层：
    #   Segment 0: [0, 1) - cond = await get_condition()
    #   Segment 1: [1, 3) - x = 0, if 语句（包含子语句块，有 await）
    #   Segment 4: [3, 4) - return x
    # then 分支：
    #   Segment 2: [0, 1) - x = await async_step1()
    # else 分支：
    #   Segment 3: [0, 1) - x = await async_step2()
    assert len(async_dataflow.segments) == 5

    assert_segment(async_dataflow.segments[0], 0, 0, 1, has_await=True)
    assert async_dataflow.segments[0].contains_sub_blocks is False

    assert_segment(async_dataflow.segments[1], 1, 1, 3, has_await=True)
    assert async_dataflow.segments[1].contains_sub_blocks is True

    assert_segment(async_dataflow.segments[2], 2, 0, 1, has_await=True)
    assert async_dataflow.segments[2].contains_sub_blocks is False

    assert_segment(async_dataflow.segments[3], 3, 0, 1, has_await=True)
    assert async_dataflow.segments[3].contains_sub_blocks is False

    assert_segment(async_dataflow.segments[4], 4, 3, 4, has_await=False)
    assert async_dataflow.segments[4].contains_sub_blocks is False

    # 验证最小代码单元划分（5 个单元）
    # 外层：
    #   Unit 0: [0, 1) - await get_condition()
    #   Unit 1: [1, 2) - cond = ..., x = 0（if 之前）
    #   Unit 4: [3, 4) - return x
    # then 分支：
    #   Unit 2: [0, 1) - await async_step1()，parent=1
    # else 分支：
    #   Unit 3: [0, 1) - await async_step2()，parent=1
    assert len(async_dataflow.units) == 5

    assert_unit(async_dataflow.units[0], 0, 0, 1, has_await=True, vars_defined=[])
    assert_unit(
        async_dataflow.units[1],
        1,
        1,
        2,
        has_await=False,
        vars_defined=["cond", "x"],
    )
    assert_unit(async_dataflow.units[2], 2, 0, 1, has_await=True, vars_defined=[])
    assert_unit(async_dataflow.units[3], 3, 0, 1, has_await=True, vars_defined=[])
    assert_unit(async_dataflow.units[4], 4, 3, 4, has_await=False, vars_defined=[])

    # 验证变量提升
    # cond: 在 Unit 1 定义，在 if 条件中使用
    #       虽然定义和使用在同一 Unit，但 if 语句包含 await
    #       co_reenter 可能跳过 cond 的初始化，所以应该是 header
    assert_var_promotion(async_dataflow, "cond", "header", defined_in_unit=1)

    # x: 在 Unit 1 定义，在 Unit 4 使用，中间有 Unit 2, 3 包含 await → member
    assert_var_promotion(async_dataflow, "x", "member", defined_in_unit=1)


def test_nested_if_with_await():
    """测试嵌套 if 中的 await"""

    code = """
async def nested_if() -> int:
    x = await async_step1()
    result = 0
    if x > 0:
        y = 0
        if x > 10:
            y = await async_step2()
        else:
            y = await async_step3()
        result = x + y
    else:
        result = 0
    return result
"""

    ir_module, async_funcs = compile_and_get_async_dataflow(code)
    assert "nested_if" in async_funcs

    func_node, async_dataflow = async_funcs["nested_if"]

    # 验证 Segment 划分（8 个 Segment）
    assert len(async_dataflow.segments) == 8

    # 验证至少有一个 Segment 包含子语句块（嵌套 if）
    has_sub_blocks = any(seg.contains_sub_blocks for seg in async_dataflow.segments)
    assert has_sub_blocks

    # 验证至少有一个 Segment 包含 await
    await_segments = [seg for seg in async_dataflow.segments if seg.has_await]
    assert len(await_segments) >= 1

    # 验证最小代码单元划分（8 个单元）
    assert len(async_dataflow.units) == 8

    # 验证变量提升
    # x: 在 Unit 1 定义，在 Unit 5 使用（result = x + y），跨越嵌套 if 中的 await → member
    assert_var_promotion(async_dataflow, "x", "member", defined_in_unit=1)

    # result: 在 Unit 1 定义，在 Unit 7 使用（return result），跨越 await → member
    assert_var_promotion(async_dataflow, "result", "member", defined_in_unit=1)

    # y: 在 Unit 2 定义，在 Unit 5 使用，跨越 await → member
    assert_var_promotion(async_dataflow, "y", "member", defined_in_unit=2)

    # 验证成员变量列表
    assert "x" in async_dataflow.member_variables
    assert "result" in async_dataflow.member_variables
    assert "y" in async_dataflow.member_variables


def test_variable_in_mixed_positions():
    """测试变量在不同分支的不同位置定义和使用"""

    code = """
async def mixed_positions(condition: bool) -> int:
    x = 0
    if condition:
        result = 10  # then 分支：await 之前定义
        y = await async_step1()
        x = result + y  # then 分支：使用 result
    else:
        result = await async_step2()  # else 分支：await 之后定义
        x = result  # else 分支：使用 result
    return x
"""

    ir_module, async_funcs = compile_and_get_async_dataflow(code)
    assert "mixed_positions" in async_funcs

    func_node, async_dataflow = async_funcs["mixed_positions"]

    # 验证 result 的提升类型
    # result 在 then 分支 await 之前定义（Unit 1），在 Unit 2 使用
    # result 在 else 分支 await 之后定义（Unit 3 的 await 结果），在 Unit 4 使用
    # 虽然 then 分支单独看可能是 header（第一个单元使用），
    # 但从 Unit 1 到 Unit 4 跨越了 await，所以整体是 member
    assert_var_promotion(async_dataflow, "result", "member")
    assert_var_promotion(async_dataflow, "x", "member")
    assert_var_promotion(async_dataflow, "y", "local")


def test_multiple_awaits_in_sequence():
    """测试连续多个 await 语句"""

    code = """
async def multiple_awaits() -> int:
    a = await step1()
    b = await step2()
    c = await step3()
    d = await step4()
    return a + b + c + d
"""

    ir_module, async_funcs = compile_and_get_async_dataflow(code)
    assert "multiple_awaits" in async_funcs

    func_node, async_dataflow = async_funcs["multiple_awaits"]

    # 验证 Segment 划分（5 个 Segment）
    # Segment 0: [0, 1) - await step1()
    # Segment 1: [1, 2) - a = ..., await step2()
    # Segment 2: [2, 3) - b = ..., await step3()
    # Segment 3: [3, 4) - c = ..., await step4()
    # Segment 4: [4, 5) - d = ..., return a + b + c + d
    assert len(async_dataflow.segments) == 5

    # 验证所有 Segment 中前 4 个包含 await
    for i in range(4):
        assert async_dataflow.segments[i].has_await

    # 验证最后一个 Segment 不包含 await
    assert not async_dataflow.segments[4].has_await

    # 验证最小代码单元划分（5 个最小代码单元）
    # Unit 0: [0, 1) - await step1()
    # Unit 1: [1, 2) - a = ..., await step2()
    # Unit 2: [2, 3) - b = ..., await step3()
    # Unit 3: [3, 4) - c = ..., await step4()
    # Unit 4: [4, 5) - d = ..., return a + b + c + d
    assert len(async_dataflow.units) == 5

    # 验证 await 结果变量的提升
    # a: 在 Unit 1 定义，在 Unit 4 使用，Unit 1 有 await → member
    assert_var_promotion(async_dataflow, "a", "member", defined_in_unit=1)

    # b: 在 Unit 2 定义，在 Unit 4 使用，Unit 2 有 await → member
    assert_var_promotion(async_dataflow, "b", "member", defined_in_unit=2)

    # c: 在 Unit 3 定义，在 Unit 4 使用，Unit 3 有 await → member
    assert_var_promotion(async_dataflow, "c", "member", defined_in_unit=3)

    # d: 在 Unit 4 定义和使用，不跨 await → local
    assert_var_promotion(async_dataflow, "d", "local", defined_in_unit=4)


def test_for_loop_without_await():
    """测试不包含 await 的 for 循环 - 循环变量应该是局部变量"""

    code = """
async def step() -> int:
    return 10

async def compute(count: int) -> int:
    a = 0
    for _ in range(count):
        a += 1
    v = await step()
    return v + 5 + a
"""

    ir_module, async_functions = compile_and_get_async_dataflow(code, add_helpers=False)

    assert "compute" in async_functions
    func_node, async_dataflow = async_functions["compute"]

    # 验证 Segment 划分（2 个 Segment）
    # 注意：Segment 包含子语句块，所以循环体不单独计算 Segment
    assert len(async_dataflow.segments) == 2

    # Segment 0: [0, 3) - 包含 a=0, for 循环, await step()
    # 注意：for 循环不包含 await，所以作为普通语句加入当前块
    # await step() 是当前块的最后一个语句
    assert_segment(async_dataflow.segments[0], 0, 0, 3, has_await=True)

    # Segment 1: [3, 4) - 包含 return v + 5 + a
    assert_segment(async_dataflow.segments[1], 1, 3, 4, has_await=False)

    # 验证最小代码单元划分（2 个 Unit）
    # 注意：不包含 await 的 for 循环不创建独立的 Unit
    # Unit 0: 外层 [0, 3) - 包含 a=0, for 循环（包括循环体 a += 1）, await step()
    # Unit 1: 外层 [3, 4) - 包含 v=..., return v + 5 + a
    assert len(async_dataflow.units) == 2

    # 按 unit_id 排序
    units_by_id = {unit.unit_id: unit for unit in async_dataflow.units}

    # Unit 0: 外层 [0, 3) - 包含 a=0, for 循环, await step()
    outer_unit_1 = units_by_id[0]
    assert_unit(outer_unit_1, 0, 0, 3, has_await=True, vars_defined=["a"])

    # Unit 1: 外层 [3, 4) - 包含 v=..., return v + 5 + a
    outer_unit_2 = units_by_id[1]
    assert_unit(outer_unit_2, 1, 3, 4, has_await=False, vars_defined=["v"])

    # 验证变量提升
    # a: 在外层 Unit 0 定义，在 Unit 1 使用，跨 await → member
    assert_var_promotion(async_dataflow, "a", "member", defined_in_unit=0)

    # v: 在外层 Unit 1 定义和使用 → local
    assert_var_promotion(async_dataflow, "v", "local", defined_in_unit=1)

    # _: 循环变量，只在循环体内使用
    # 预期行为：循环体中没有使用 _，所以 _ 不应该在 variable_promotions 中
    # _ 不应该被提升为 member 变量或 header 变量
    # 如果将来循环体中使用了 _，那么 _ 应该在 variable_promotions 中，且应为 local 类型
    assert "_" not in async_dataflow.variable_promotions
    assert "_" not in async_dataflow.member_variables
    assert "_" not in async_dataflow.header_variables


def test_if_branches_with_await():
    """测试 if/else 分支中都有 await 的情况"""

    code = """
async def get_value(x: int) -> int:
    return x

async def process(use_then: bool) -> int:
    if use_then:
        x = await get_value(10)
        return x
    else:
        y = await get_value(20)
        return y
"""

    ir_module, async_functions = compile_and_get_async_dataflow(code, add_helpers=False)

    # 获取 process 函数的数据流分析结果
    assert "process" in async_functions
    func_node, async_dataflow = async_functions["process"]

    # 验证变量提升
    # x: 在 then 分支的 await 之后定义，在同一分支使用 → 应该是 local
    # y: 在 else 分支的 await 之后定义，在同一分支使用 → 应该是 local
    assert_var_promotion(async_dataflow, "x", "local")
    assert_var_promotion(async_dataflow, "y", "local")

    # 验证变量不在 member_variables 或 header_variables 中
    assert "x" not in async_dataflow.member_variables
    assert "x" not in async_dataflow.header_variables
    assert "y" not in async_dataflow.member_variables
    assert "y" not in async_dataflow.header_variables


def test_await_followed_by_if_without_await():
    """测试 await 后紧跟不包含 await 的 if 语句

    关键点：
    1. await 结果变量在 await 语句中定义
    2. 紧跟的不包含 await 的 if 语句不应该中断 Unit
    3. await 结果变量应该在同一 Unit 中定义和使用，不应该被提升为 member
    """

    code = """
async def get_value(x: int) -> int:
    return x + 10

async def process() -> int:
    total = 0
    for i in range(3):
        val = await get_value(i)
        if val % 2 == 0:
            continue
        total = total + val
    return total
"""

    ir_module, async_functions = compile_and_get_async_dataflow(code, add_helpers=False)

    assert "process" in async_functions
    func_node, async_dataflow = async_functions["process"]

    # 验证 Segment 划分
    # 外层：
    #   Segment 0: [0, 2) - total = 0, for 循环（包含子语句块，有 await）
    #   Segment 3: [2, 3) - return total
    # 循环体：
    #   Segment 1: [0, 1) - await get_value(i)
    #   Segment 2: [1, 3) - val = ..., if val % 2 == 0: continue, total = total + val
    #   注意：不包含 await 的 if 语句不会创建新的 Segment
    assert len(async_dataflow.segments) == 4

    assert_segment(async_dataflow.segments[0], 0, 0, 2, has_await=True)
    assert async_dataflow.segments[0].contains_sub_blocks is True

    assert_segment(async_dataflow.segments[1], 1, 0, 1, has_await=True)
    assert async_dataflow.segments[1].contains_sub_blocks is False

    assert_segment(async_dataflow.segments[2], 2, 1, 3, has_await=False)
    assert async_dataflow.segments[2].contains_sub_blocks is False

    assert_segment(async_dataflow.segments[3], 3, 2, 3, has_await=False)
    assert async_dataflow.segments[3].contains_sub_blocks is False

    # 验证最小代码单元划分
    # 关键：不包含 await 的 if 语句不应该中断 Unit
    # 外层：
    #   Unit 0: [0, 1) - total = 0（for 之前）
    #   Unit 3: [2, 3) - return total（for 之后）
    # 循环体：
    #   Unit 1: [0, 1) - await get_value(i)，segment=1
    #   Unit 2: [1, 3) - val = ..., if val % 2 == 0: continue, total = total + val，segment=1
    #   注意：val 的定义和使用都在 Unit 2 中
    assert len(async_dataflow.units) == 4

    assert_unit(
        async_dataflow.units[0],
        0,
        0,
        1,
        has_await=False,
        vars_defined=["total"],
    )
    assert_unit(async_dataflow.units[1], 1, 0, 1, has_await=True, vars_defined=[])
    assert_unit(
        async_dataflow.units[2],
        2,
        1,
        3,
        has_await=False,
        vars_defined=["val"],  # val 是 await 结果变量，在此 Unit 中定义
    )
    assert_unit(async_dataflow.units[3], 3, 2, 3, has_await=False, vars_defined=[])

    # 验证变量提升
    # total: 在 Unit 0 定义，在 Unit 2, 3 使用，跨越 Unit 1 的 await → member
    assert_var_promotion(async_dataflow, "total", "member", defined_in_unit=0)

    # val: 在 Unit 2 定义和使用（if 条件和 total = total + val）→ local
    # 关键验证点：val 不应该被提升为 member 或 header
    assert_var_promotion(async_dataflow, "val", "local", defined_in_unit=2)
    assert "val" not in async_dataflow.member_variables
    assert "val" not in async_dataflow.header_variables

    # i: 循环变量，在循环体中使用（Unit 1: get_value(i)），跨 await → member
    assert_var_promotion(async_dataflow, "i", "member")


def test_await_followed_by_if_with_break():
    """测试 await 后紧跟不包含 await 的 if 语句（带 break）

    与 continue 类似，break 语句也不应该影响 Unit 划分
    """

    code = """
async def get_value(x: int) -> int:
    return x + 10

async def process() -> int:
    total = 0
    for i in range(10):
        val = await get_value(i)
        if val > 15:
            break
        total = total + val
    return total
"""

    ir_module, async_functions = compile_and_get_async_dataflow(code, add_helpers=False)

    assert "process" in async_functions
    func_node, async_dataflow = async_functions["process"]

    # 验证 Segment 划分（4 个 Segment）
    assert len(async_dataflow.segments) == 4

    # 验证最小代码单元划分（4 个 Unit）
    # val 的定义和使用应该在同一个 Unit 中
    assert len(async_dataflow.units) == 4

    # 找到循环体中 await 后的 Unit（segment_id=1 表示循环体内的 Segment）
    loop_units = [
        u for u in async_dataflow.units if u.segment_id == 1 and not u.has_await
    ]
    assert len(loop_units) == 1
    loop_unit = loop_units[0]

    # 验证 val 在此 Unit 中定义
    assert "val" in loop_unit.variables_defined

    # 验证变量提升
    # val: 在循环体的 await 后 Unit 中定义和使用 → local
    assert_var_promotion(async_dataflow, "val", "local")
    assert "val" not in async_dataflow.member_variables
    assert "val" not in async_dataflow.header_variables


def test_result_variable_in_if_after_await():
    """测试在 await 之后的 if/else 中定义 result 变量

    这是一个关键测试：验证在 await 之后的 if/else 分支中定义的变量，
    如果只在最后一个 Unit 中使用，不应该被提升为 member 变量。
    """

    code = """
async def async_step(x: int) -> int:
    return x * 2

async def compute() -> int:
    a = 0
    b = 0
    for i in range(3):
        a = a + i
    v1 = await async_step(a)
    
    if v1 > 5:
        b = 10
    else:
        b = 20
    
    c = 0
    for i in range(2):
        for j in range(2):
            c = c + i + j
    v2 = await async_step(c)
    
    if v2 > 10:
        result = v1 + v2 + b
    else:
        result = v1 + v2 - b
    
    return result
"""

    ir_module, async_funcs = compile_and_get_async_dataflow(code, add_helpers=False)
    assert "compute" in async_funcs

    func_node, async_dataflow = async_funcs["compute"]

    # 验证 Unit 数量：两个 await 应该产生 3 个 Unit
    assert len(async_dataflow.units) == 3

    # 验证 result 在最后一个 Unit (Unit 2) 中定义
    assert "result" in async_dataflow.units[2].variables_defined

    # 验证变量提升：result 应该在最后一个 Unit 中定义和使用，不跨 await → local
    assert_var_promotion(async_dataflow, "result", "local", defined_in_unit=2)
    assert "result" not in async_dataflow.member_variables
    assert "result" not in async_dataflow.header_variables


# def test_for_each_loop_with_await():
#     """测试 for-each 循环中包含 await 的情况 - 循环变量应该被提升为成员变量"""

#     code = """
# async def get_value(x: int) -> int:
#     return x + 10

# async def process() -> int:
#     total = 0
#     numbers = [1, 2, 3, 4, 5]
#     for num in numbers:
#         val = await get_value(num)
#         total = total + val
#     return total
# """

#     ir_module, async_functions = compile_and_get_async_dataflow(code, add_helpers=False)

#     assert "process" in async_functions
#     func_node, async_dataflow = async_functions["process"]

#     # 验证 Segment 划分（4 个 Segment）
#     # 外层：
#     #   Segment 0: [0, 3) - total = 0, numbers = [1, 2, 3, 4, 5], for 循环（包含子语句块，有 await）
#     #   Segment 3: [3, 4) - return total
#     # 循环体：
#     #   Segment 1: [0, 1) - await get_value(num)
#     #   Segment 2: [1, 2) - val = ..., total = total + val
#     assert len(async_dataflow.segments) == 4

#     assert_segment(async_dataflow.segments[0], 0, 0, 3, has_await=True)
#     assert async_dataflow.segments[0].contains_sub_blocks is True

#     assert_segment(async_dataflow.segments[1], 1, 0, 1, has_await=True)
#     assert async_dataflow.segments[1].contains_sub_blocks is False

#     assert_segment(async_dataflow.segments[2], 2, 1, 2, has_await=False)
#     assert async_dataflow.segments[2].contains_sub_blocks is False

#     assert_segment(async_dataflow.segments[3], 3, 3, 4, has_await=False)
#     assert async_dataflow.segments[3].contains_sub_blocks is False

#     # 验证最小代码单元划分（4 个单元）
#     # 外层：
#     #   Unit 0: [0, 2) - total = 0, numbers = [1, 2, 3, 4, 5]（for 之前）
#     #   Unit 3: [3, 4) - return total
#     # 循环体：
#     #   Unit 1: [0, 1) - await get_value(num)，parent=0
#     #   Unit 2: [1, 2) - val = ..., total = total + val，parent=0
#     assert len(async_dataflow.units) == 4

#     assert_unit(
#         async_dataflow.units[0],
#         0,
#         0,
#         2,
#         has_await=False,
#         vars_defined=["total", "numbers"],
#     )
#     assert_unit(async_dataflow.units[1], 1, 0, 1, has_await=True, vars_defined=[])
#     assert_unit(
#         async_dataflow.units[2],
#         2,
#         1,
#         2,
#         has_await=False,
#         vars_defined=["val"],
#     )
#     assert_unit(async_dataflow.units[3], 3, 3, 4, has_await=False, vars_defined=[])

#     # 验证变量提升
#     # total: 在 Unit 0 定义，在 Unit 2, 3 使用，跨越 Unit 1 的 await → member
#     assert_var_promotion(async_dataflow, "total", "member", defined_in_unit=0)

#     # val: 在 Unit 2 定义和使用 → local
#     assert_var_promotion(async_dataflow, "val", "local", defined_in_unit=2)

#     # num: 循环变量，由 for 语句隐式定义
#     #      在循环体第一个 Unit 中使用（Unit 1: get_value(num)）
#     #      根据规则：
#     #      1. 不属于同一个 unit 定义同一个 unit 使用，不属于 local
#     #      2. 不属于跨过 await 使用（在 await 之前使用），不属于 member
#     #      3. 属于其他类型 header 变量
#     assert_var_promotion(async_dataflow, "num", "header")


def test_iterator_variable_with_await():
    """测试使用 iter() 和 next() 的迭代器变量 - 应该被提升为成员变量

    关键点：
    1. it = iter(numbers) 创建迭代器
    2. it 在 while 循环的 next(it, None) 中使用，而 next() 调用在包含 await 的 co_yield 块内
    3. it 跨越 await 边界使用，应该被提升为 member 变量
    """

    code = """
async def get_value(x: int) -> int:
    return x + 10

async def process() -> int:
    total = 0
    numbers = [1, 2, 3, 4, 5]
    it = iter(numbers)
    while True:
        num = next(it, None)
        if not num:
            break
        val = await get_value(num)
        total = total + val
    return total
"""

    _, async_functions = compile_and_get_async_dataflow(code, add_helpers=False)

    assert "process" in async_functions
    _, async_dataflow = async_functions["process"]

    # 验证变量提升
    # total: 在循环外定义，在循环内和循环后使用，跨 await → member
    assert_var_promotion(async_dataflow, "total", "member")

    # it: 关键测试点！
    # it 在循环外定义（it = iter(numbers)）
    # it 在循环内使用（next(it, None)），而这个使用在包含 await 的块内
    # it 跨越 await 边界 → 应该是 member 变量
    assert_var_promotion(async_dataflow, "it", "member")


def test_dict_comp_loop_variable_promotion():
    """测试字典推导式中循环变量的提升 - 关键测试用例

    字典推导式: {x: await get_value(x) for x in numbers}
    等价于手动展开：
        result = {}
        it = iter(numbers)
        while True:
            x = next(it, None)
            if not x:
                break
            value = await get_value(x)
            result[x] = value  # x 在 await 之后使用！

    关键点：
    1. 循环变量 x 在 await 之前定义（x = next(it, None)）
    2. x 在 await 之后使用（result[x] = value）
    3. x 跨越 await 边界 → 应该是 member 变量

    这个问题与列表推导式带条件的区别：
    - 列表推导式: [await get_value(x) for x in numbers if x % 2 == 0]
      x 在条件中使用（await 之前），在 await 之后不再使用
    - 字典推导式: {x: await get_value(x) for x in numbers}
      x 作为 key 使用（await 之后），所以必须被提升
    """

    code = """
async def get_value(x: int) -> int:
    return x * 2

async def process() -> dict:
    numbers = [1, 2, 3, 4, 5]
    result = {}
    it = iter(numbers)
    while True:
        x = next(it, None)
        if not x:
            break
        value = await get_value(x)
        result[x] = value
    return result
"""

    _, async_functions = compile_and_get_async_dataflow(code, add_helpers=False)

    assert "process" in async_functions
    _, async_dataflow = async_functions["process"]

    # 验证变量提升
    # x: 关键测试点！
    # x 在循环内定义（x = next(it, None)），在 await 之后使用（result[x] = value）
    # x 跨越 await 边界（在 Unit 1 和 Unit 2 中使用）→ 应该是 member 变量
    assert_var_promotion(async_dataflow, "x", "member", defined_in_unit=1)

    # 验证其他变量
    assert_var_promotion(async_dataflow, "result", "member")
    assert_var_promotion(async_dataflow, "it", "member")
    assert_var_promotion(async_dataflow, "value", "local")
